# -*- coding: utf-8 -*-

"""
    将原始数据集进行划分成训练集、验证集和测试集
"""

import os
import glob
import random
import shutil
from utils.utils import get_classes
from utils.utils import makedir

#--------------------------------------------------------------------------------------------------------------------------------#
#   annotation_mode用于指定该文件运行时计算的内容
#   annotation_mode为0代表整个数据集处理过程,包括获得datasets里面的train、test、valid文件夹及下边的图片,以及训练用的cls_train.txt、cls_valid.txt
#   annotation_mode为1代表获得datasets里面的train、test、valid文件夹及下边的图片
#   annotation_mode为2代表获得训练用的cls_train.txt、cls_valid.txt、cls_test.txt
#--------------------------------------------------------------------------------------------------------------------------------#
annotation_mode     = 0

#--------------------------------------------#
#   train_per, valid_per, test_per分别用于
#   指定数据集中训练集, 验证集和测试集的比例
#   默认情况下比例为8:1:1
#--------------------------------------------#
train_per    = 0.7
valid_per    = 0.3
test_per     = 0

#-------------------------------------------------------------------#
#   classes_path    指向model_data下的txt，与自己训练的数据集相关
#                   训练前一定要修改classes_path，使其对应自己的数据集
#                   txt文件中是自己所要去区分的种类
#                   与训练和预测所用的classes_path一致即可
#-------------------------------------------------------------------#
classes_path    = 'model_data/news_classes.txt'

#-------------------------------------------------------#
#   datasets_path   指向数据集所在的文件夹
#-------------------------------------------------------#
dataset_dir = "datasets"

train_dir = os.path.join(".", dataset_dir, "train")
valid_dir = os.path.join(".", dataset_dir, "valid")
test_dir = os.path.join(".", dataset_dir, "test")

train_txt_path = os.path.join(".", "cls_train.txt")
valid_txt_path = os.path.join(".", "cls_valid.txt")
test_txt_path = os.path.join(".", "cls_test.txt")

classes, _      = get_classes(classes_path)


# 生成包含图片路径和对应类的txt文件
def gen_txt(txt_path, img_dir):
    f = open(txt_path, 'w')
    for root, s_dirs, _ in os.walk(img_dir, topdown=True):  # 获取 train文件下各文件夹名称
        for sub_dir in s_dirs:
            i_dir = os.path.join(root, sub_dir)             # 获取各类的文件夹 绝对路径
            img_list = os.listdir(i_dir)                    # 获取类别文件夹下所有jpg图片的路径
            for i in range(len(img_list)):
                # 若不是jpg、jpeg、png文件，跳过
                if img_list[i].endswith("jpg") or img_list[i].endswith("jpeg") or img_list[i].endswith("png"):
                    cls_id = classes.index(sub_dir)
                    label = str(cls_id)
                    img_path = os.path.join(i_dir, img_list[i])
                    line = img_path + ' ' + label + '\n'
                    f.write(line)
                else:
                    continue
    f.close()


if __name__ == '__main__':

    if annotation_mode == 0 or annotation_mode == 1:
        # --------------------------------------------#
        #   划分训练集, 验证集和测试集
        #   生成train, valid, test三个文件夹
        #   及文件夹下的图片
        # --------------------------------------------#
        print("Generate train, valid, test folders in %s." % dataset_dir)

        #  使用os.walk来遍历文件夹中的文件
        #  root 表示当前正在访问的文件夹路径
        #  dirs 表示该文件夹下的子目录名list
        #  files 表示该文件夹下的文件list
        for root, dirs, files in os.walk(dataset_dir):
            for sDir in dirs:
                imgs_list = glob.glob(os.path.join(root, sDir, '*.jpg'))
                random.seed(666)
                random.shuffle(imgs_list)
                imgs_num = len(imgs_list)

                train_point = int(imgs_num * train_per)
                valid_point = int(imgs_num * (train_per + valid_per))

                for i in range(imgs_num):
                    if i < train_point:
                        out_dir = os.path.join(train_dir, sDir)
                    elif i < valid_point:
                        out_dir = os.path.join(valid_dir, sDir)
                    else:
                        out_dir = os.path.join(test_dir, sDir)

                    makedir(out_dir)
                    out_path = os.path.join(out_dir, os.path.split(imgs_list[i])[-1])
                    shutil.copy(imgs_list[i], out_path)

                print('Class:{}, train:{}, valid:{}, test:{}'.format(sDir, train_point, valid_point-train_point, imgs_num-valid_point))

    if annotation_mode == 0 or annotation_mode == 2:

        # ---------------------------------- #
        #   生成train, valid, test三个txt文件
        # ---------------------------------- #
        print("Generate txt in %s." % dataset_dir)
        gen_txt(train_txt_path, train_dir)
        gen_txt(valid_txt_path, valid_dir)
        gen_txt(test_txt_path, test_dir)
        print("Generate txt done.")


